from .ConvLSTM import ConvLSTM
from .Encode import Encode
from .Res_MCNN import Res_MCNN
from .GCN import GCN
import torch
import torch.nn as nn
class RF_STED_No_S(nn.Module):
    def __init__(self):
        super(RF_STED_No_S,self).__init__()
        self.covlstm = ConvLSTM(input_dim = 1, 
                                hidden_dim = 24, 
                                kernel_size = (3,3),
                                num_layers = 1,
                                batch_first=True, 
                                bias=True, 
                                return_all_layers=False)
        self.encode = Encode()
        self.Res_CNN1 = Res_MCNN(num_convolutions=3, kernel_size=31, scale=2)
        self.Res_CNN2 = Res_MCNN(num_convolutions=3, kernel_size=33, scale=2)
        self.Res_CNN3 = Res_MCNN(num_convolutions=3, kernel_size=35, scale=2)   #num_convolutions越小越好  kernel_size居中  scale 2 差不多
        self.relu = nn.ReLU(inplace=False) #relu激活
        self.sig = nn.Sigmoid()
    def forward(self,x,adj1,adj2,adj3,adj4):  #前向传播   x:[6,1,6,69,69]  adj:[4761,4761]
        #时间
        taxi_three_time_x = x[:,0,0:3,:,:]   #[6, 3, 69, 69]
        taxi_three_day_x = x[:,0,3:6,:,:]    #[6, 3, 69, 69]
        taxi_three_time_x = torch.unsqueeze(taxi_three_time_x,2) #[6, 3, 1,69, 69]
        taxi_three_day_x = torch.unsqueeze(taxi_three_day_x,2)   #[6, 3, 1,69, 69]
        taxi_three_time_x = self.covlstm(taxi_three_time_x)[1][0][0]  #[6,32,69,69]
        taxi_three_day_x = self.covlstm(taxi_three_day_x)[1][0][0]   #[6,32,69,69]
        taxi_three_time_x = taxi_three_time_x.permute(0,2,3,1)  #[6,69,69,32]
        taxi_three_day_x = taxi_three_day_x.permute(0,2,3,1)   #[6,69,69,32]
        # 编码
        time = torch.cat((taxi_three_time_x.unsqueeze(-1), taxi_three_day_x.unsqueeze(-1)), dim=-1)         #[6,69,69,32,2]
        x1_merged_tensor = time.view(time.shape[0], time.shape[1], time.shape[2] , time.shape[3] * time.shape[4])  #[6,69,69,64]
        encoder = x1_merged_tensor.sum(dim=3, keepdim=True)  #[6,69,69,1]
        # 解码
        encoder = encoder.permute(0,3,1,2)   #[6,1,69,69]
        encoder = self.Res_CNN1(encoder) #[6,1,69,69]
        encoder = self.Res_CNN2(encoder) #[6,1,69,69]
        encoder = self.Res_CNN3(encoder) #[6,1,69,69]
        out = self.relu(encoder)
        return out
